# Modify Date: 2021/6/22
# Usage: output connected line include row and col
# Modify Date: 2021/8/13
# Usage: Connect gap from the same line
import cv2
import numpy as np
import os
import math

def line_row_gen(pre_map):

    img_temp = np.ones_like(pre_map) *255

    # gray = cv2.cvtColor( pre_map,cv2.COLOR_BGR2GRAY )
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 1))
    gray = cv2.morphologyEx(pre_map, cv2.MORPH_ERODE, kernel, iterations=10)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (6, 1))
    gray = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel, iterations=10)
    ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)

    # 检测图中直线
    # minLineLength = 100
    # maxLineGap = 100
    # lines = cv2.HoughLinesP( binary,1,np.pi/180,100,minLineLength=minLineLength,maxLineGap=maxLineGap )
    # print('lines:', len(lines))
    # try:
    #     for line in lines:
    #         for x1,y1,x2,y2 in line:
    #             # cv2.line( img,( x1,y1 ),( x2,y2 ),( 0,255,0 ),2 )
    #             cv2.line(img_temp, (x1, y1), (x2, y2), (0,), 1)
    # except:
    #     return img_temp, 0
    # cv2.imwrite( 'E:/image/myhoughlinesp.jpg',img )

    # cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    # cv2.imshow( img_path,closed_img_temp)
    # cv2.imshow(img_path, closed_img_temp)
    # cv2.waitKey(0)
    return binary

def angle(v1, v2):
  dx1 = v1[2] - v1[0]
  dy1 = v1[3] - v1[1]
  dx2 = v2[2] - v2[0]
  dy2 = v2[3] - v2[1]
  angle1 = math.atan2(dy1, dx1)
  angle1 = int(angle1 * 180/math.pi)
  angle2 = math.atan2(dy2, dx2)
  angle2 = int(angle2 * 180/math.pi)
  if angle1*angle2 >= 0:
    included_angle = abs(angle1-angle2)
  else:
    included_angle = abs(angle1) + abs(angle2)
    if included_angle > 180:
      included_angle = 360 - included_angle
  return included_angle

def line_col_gen(pre_map):
    AB = [0,0,100,0]

    img_temp = np.ones_like(pre_map) *255

    # gray = cv2.cvtColor( pre_map,cv2.COLOR_BGR2GRAY )
    ret, binary = cv2.threshold(pre_map, 127, 255, cv2.THRESH_BINARY_INV)

    # cv2.imshow('edge', edges)
    # cv2.imshow('binary',binary)
    # cv2.waitKey(0)

    # minLineLength = 100
    # maxLineGap = 100
    # lines = cv2.HoughLinesP( binary,1,np.pi/180,100,minLineLength=20,maxLineGap=50 )
    # try:
    #     for line in lines:
    #         for x1,y1,x2,y2 in line:
    #             CD = [x1, y1, x2, y2]
    #             angle_cross = angle(AB, CD)
    #             if angle_cross<90+15 and angle_cross>90-15:
    #                 cv2.line( pre_map,( x1,y1 ),( x2,y2 ),( 255,0,0 ),2 )
    #                 # cv2.line(img_temp, (x1, y1), (x2, y2), (0,), 2)
    # except:
    #     return img_temp
    # points = [(box[0], box[1]), (box[2],box[1]), (box[2], box[3]), (box[0], box[3])]
    # cv2.fillPoly(image,[np.array(points)],(255,0,0))
    return binary

def tx_post(pre_row, pre_nrow, pre_col, pre_ncol):
    post_row = line_row_gen(pre_row)
    post_nrow = line_row_gen(pre_nrow)
    post_col = line_col_gen(pre_col)
    post_ncol = line_col_gen(pre_ncol)
    return post_nrow, post_ncol, post_row, post_col


if __name__ == '__main__':
    col_root = r'.\tx_infer_data\col'
    row_root = r'.\tx_infer_data\row'
    ncol_root = r'.\tx_infer_data\ncol'
    nrow_root = r'.\tx_infer_data\nrow'
    save_col_root = r'.\tx_process_data\col'
    save_row_root = r'.\tx_process_data\row'
    save_ncol_root = r'.\tx_process_data\ncol'
    save_nrow_root = r'.\tx_process_data\nrow'

    img_names = os.listdir(col_root)
    for img_name in img_names:
        col_path = os.path.join(col_root, img_name)
        ncol_path = os.path.join(ncol_root, img_name)
        row_path = os.path.join(row_root, img_name)
        nrow_path = os.path.join(nrow_root, img_name)
        save_col_path = os.path.join(save_col_root, img_name)
        save_row_path = os.path.join(save_row_root, img_name)
        save_ncol_path = os.path.join(save_ncol_root, img_name)
        save_nrow_path = os.path.join(save_nrow_root, img_name)
        tx_post(row_path, nrow_path, col_path, ncol_path, save_row_path, save_nrow_path, save_col_path, save_ncol_path)
